package DataStructure.uf;

public class UF_Tree_Weighted {
    private int[] eleAndGroup;//记录节点元素和所在组的标识
    private int count;//组的个数
    private int[] sz;//存储每一个根节点对应的树中保存的节点的个数

    public UF_Tree_Weighted(int N){
        this.count=N;
        this.eleAndGroup=new int[N];

        for (int i = 0; i < eleAndGroup.length; i++) {
            eleAndGroup[i]=i;
        }

        //默认情况下sz中每个索引处的值都是1
        this.sz=new int[N];
        for (int i = 0; i < sz.length; i++) {
            sz[i]=1;
        }
    }

    public int count(){
        return count;
    }

    public boolean connected(int p,int q){
        return find(p)==find(q);
    }

    //元素p所在的标识符
    public int find(int p){
        while (true){
            if(p==eleAndGroup[p]){
                return p;
            }
            p=eleAndGroup[p];
        }
    }

    //合并
    public void union(int p,int q){
        //找到p和q所在组对应的数的根节点
        int pRoot = find(p);
        int qRoot = find(q);
        //如果p和q已经在同一分组就不需要合并
        if(pRoot==qRoot){
            return;
        }


        //判断proot对应的树大还是qroot对应的树大，最终把较小的树合并到较大的树中
        if(sz[pRoot]<sz[qRoot]){
            eleAndGroup[pRoot]=qRoot;
            sz[qRoot]+=sz[pRoot];
        }else {
            eleAndGroup[qRoot]=pRoot;
            sz[pRoot]+=sz[qRoot];
        }

        this.count--;
    }
}
